import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, hyper_params):
        super(Model, self).__init__()
        self.hyper_params = hyper_params

        self.dropout = nn.Dropout(hyper_params['dropout'])
        self.encoder1 = nn.Sequential(
            nn.Linear(hyper_params['device_cnt'], 16),
            nn.Sigmoid()
        )
        self.encoder2 = nn.Sequential(
            nn.Linear(16, 8),
            nn.Sigmoid()
        )

        self.decoder1 = nn.Sequential(
            nn.Linear(8, 16),
            nn.Sigmoid()
        )
        self.decoder2 = nn.Sequential(
            nn.Linear(16, hyper_params['device_cnt']),
            nn.Sigmoid()
        )

        self.pred1 = nn.Sequential(
            nn.Linear(8, 32),
            nn.ReLU()
        )
        self.pred2 = nn.Sequential(
            nn.Linear(32, 64),
            nn.ReLU()
        )
        self.pred3 = nn.Linear(64, hyper_params['max_x'] * hyper_params['max_y'])
        # self.pred2 = nn.Linear(100, 2)

    def dae(self, x):
        x = self.dropout(x)
        z = self.encoder2(self.encoder1(x))
        z1 = self.dropout(z)

        y = self.decoder2(self.decoder1(z1))
        return y, z

    def predict(self, z):
        z = self.dropout(z)
        pred = self.pred3(self.pred2(self.pred1(z)))

        return pred

class GRUPredictor(nn.Module):
    def __init__(self, hyper_params):
        super(GRUPredictor, self).__init__()

        self.hyper_params = hyper_params
        self.encoder = nn.Sequential(
            nn.Linear(hyper_params['device_cnt'], 32),
            nn.Sigmoid()
        )
        self.gru = nn.GRU(
                    input_size=32,
                    hidden_size=16,
                    num_layers=1
                    )
        self.decoder1 = nn.Sequential(
            nn.Linear(16, 32),
            nn.Sigmoid()
        )
        self.decoder2 = nn.Sequential(
            nn.Linear(32, hyper_params['device_cnt']),
            nn.Sigmoid()
        )
        self.pred1 = nn.Sequential(
            nn.Linear(16, 32),
            nn.ReLU()
        )
        self.pred2 = nn.Linear(32, hyper_params['max_x'] * hyper_params['max_y'])
        self.dropout = nn.Dropout(hyper_params['dropout'])
            
    def dae(self, x):
        x = self.dropout(x)
        z = torch.sigmoid(self.gru(self.encoder(x))[0])
        z1 = self.dropout(z)

        y = self.decoder2(self.decoder1(z1))
        return y, z1

    def predict(self, z):
        z = self.dropout(z)
        out = self.pred2(self.pred1(z))

        return out
